import RPi.GPIO as GPIO
import time
import cv2
import numpy as np

# 初始化GPIO
GPIO.setmode(GPIO.BCM)

# 定义电机控制引脚
# 假设使用两个电机：左电机和右电机
# 每个电机需要两个引脚：方向和PWM
LEFT_FORWARD = 17
LEFT_BACKWARD = 27
LEFT_PWM = 22

RIGHT_FORWARD = 23
RIGHT_BACKWARD = 24
RIGHT_PWM = 25

# 定义机械臂控制引脚（假设使用舵机）
GRIPPER_PIN = 18
ARM_PIN = 19

# 设置引脚为输出
GPIO.setup(LEFT_FORWARD, GPIO.OUT)
GPIO.setup(LEFT_BACKWARD, GPIO.OUT)
GPIO.setup(LEFT_PWM, GPIO.OUT)

GPIO.setup(RIGHT_FORWARD, GPIO.OUT)
GPIO.setup(RIGHT_BACKWARD, GPIO.OUT)
GPIO.setup(RIGHT_PWM, GPIO.OUT)

GPIO.setup(GRIPPER_PIN, GPIO.OUT)
GPIO.setup(ARM_PIN, GPIO.OUT)

# 设置PWM频率
left_pwm = GPIO.PWM(LEFT_PWM, 1000)
right_pwm = GPIO.PWM(RIGHT_PWM, 1000)
left_pwm.start(0)
right_pwm.start(0)

# 设置舵机PWM
gripper_pwm = GPIO.PWM(GRIPPER_PIN, 50)  # 50Hz
arm_pwm = GPIO.PWM(ARM_PIN, 50)
gripper_pwm.start(7.5)  # 中间位置
arm_pwm.start(7.5)

def move_forward(speed=50):
    left_pwm.ChangeDutyCycle(speed)
    right_pwm.ChangeDutyCycle(speed)
    GPIO.output(LEFT_FORWARD, GPIO.HIGH)
    GPIO.output(LEFT_BACKWARD, GPIO.LOW)
    GPIO.output(RIGHT_FORWARD, GPIO.HIGH)
    GPIO.output(RIGHT_BACKWARD, GPIO.LOW)

def move_backward(speed=50):
    left_pwm.ChangeDutyCycle(speed)
    right_pwm.ChangeDutyCycle(speed)
    GPIO.output(LEFT_FORWARD, GPIO.LOW)
    GPIO.output(LEFT_BACKWARD, GPIO.HIGH)
    GPIO.output(RIGHT_FORWARD, GPIO.LOW)
    GPIO.output(RIGHT_BACKWARD, GPIO.HIGH)

def turn_left(speed=50):
    left_pwm.ChangeDutyCycle(speed)
    right_pwm.ChangeDutyCycle(speed)
    GPIO.output(LEFT_FORWARD, GPIO.LOW)
    GPIO.output(LEFT_BACKWARD, GPIO.HIGH)
    GPIO.output(RIGHT_FORWARD, GPIO.HIGH)
    GPIO.output(RIGHT_BACKWARD, GPIO.LOW)

def turn_right(speed=50):
    left_pwm.ChangeDutyCycle(speed)
    right_pwm.ChangeDutyCycle(speed)
    GPIO.output(LEFT_FORWARD, GPIO.HIGH)
    GPIO.output(LEFT_BACKWARD, GPIO.LOW)
    GPIO.output(RIGHT_FORWARD, GPIO.LOW)
    GPIO.output(RIGHT_BACKWARD, GPIO.HIGH)

def stop():
    left_pwm.ChangeDutyCycle(0)
    right_pwm.ChangeDutyCycle(0)
    GPIO.output(LEFT_FORWARD, GPIO.LOW)
    GPIO.output(LEFT_BACKWARD, GPIO.LOW)
    GPIO.output(RIGHT_FORWARD, GPIO.LOW)
    GPIO.output(RIGHT_BACKWARD, GPIO.LOW)

def set_gripper(angle):
    duty = angle / 180 * 12 + 2  # 简单的线性映射
    gripper_pwm.ChangeDutyCycle(duty)

def set_arm(angle):
    duty = angle / 180 * 12 + 2
    arm_pwm.ChangeDutyCycle(duty)

def detect_object():
    # 使用OpenCV进行简单的颜色检测
    cap = cv2.VideoCapture(0)
    while True:
        ret, frame = cap.read()
        hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
        # 定义颜色范围（红色）
        lower_red = np.array([0, 120, 70])
        upper_red = np.array([10, 255, 255])
        mask1 = cv2.inRange(hsv, lower_red, upper_red)
        lower_red = np.array([170, 120, 70])
        upper_red = np.array([180, 255, 255])
        mask2 = cv2.inRange(hsv, lower_red, upper_red)
        mask = mask1 | mask2
        contours, _ = cv2.findContours(mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
        for cnt in contours:
            area = cv2.contourArea(cnt)
            if area > 500:
                x, y, w, h = cv2.boundingRect(cnt)
                cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 255, 0), 2)
                return x, y, w, h
        cv2.imshow('frame', frame)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
    cap.release()
    cv2.destroyAllWindows()

try:
    while True:
        # 简单的控制逻辑
        x, y, w, h = detect_object()
        # 移动到物体位置
        if x > 320:
            turn_right(50)
        elif x < 320:
            turn_left(50)
        else:
            move_forward(50)
        # 停止并抓取
        stop()
        set_gripper(30)  # 打开抓取器
        time.sleep(1)
        set_gripper(0)   # 关闭抓取器
        time.sleep(1)
        # 移动到目标位置
        move_backward(50)
        time.sleep(2)
        stop()
        break
except KeyboardInterrupt:
    print("程序终止")
finally:
    GPIO.cleanup()